-- [[ uncomment it to disable this module ]]--
-- do return end

g_rtu = Proto("ModbusRTU", "Modbus RTU")

print("Script ModbusRTU", NotWireshark)

function tableContains(tbl, key)
    if tbl == nil then
        return false
    end
    for k, v in pairs(tbl) do
        if k == key then
            return true
        end
    end

    return false
end

-- g_rtu.enable = false

local functions_code = {
    [1] = "Read Coils",
    [2] = "Read Discrete Inputs",
    [3] = "Read Holding Registers",
    [4] = "Read Input Registers",
    [5] = "Write Single Coil",
    [6] = "Write Holding Register",
    [0x0F] = "Write (N) Coils",
    [0x10] = "Write (N) Holding Registers",

    [0x81] = "Read Coils error",
    [0x82] = "Read Discrete Inputs error",
    [0x83] = "Read Holding Registers error",
    [0x84] = "Read Input Registers error",
    [0x85] = "Write Single Coil error",
    [0x86] = "Write Holding Register error",
    [0x8F] = "Write (N) Coils error",
    [0x90] = "Write (N) Holding Registers error",
}

local exception_code = {
    [1] = "Illegal function",
    [2] = "Illegal address",
    [3] = "Illegal data",
    [4] = "Slave device failure",
    [5] = "Ack",
    [6] = "Slave device busy",
    [8] = "Checksum error",
    [0xA] = "Gateway unavailable",
    [0XB] = "Gateway response failure"
}

local coils_code = {
    [0] = "Off",
    [0x00FF] = "On",
    [0xFF00] = "On"
}

-- define protocol fields
-- local f1 = ProtoField.new("rtu.BOOLEAN", "BOOLEAN", ftypes.BOOLEAN)
-- local pf_bool = ProtoField.bool("rtu.BOOLEAN", "BOOLEAN", base.NONE)--, base.HEX, {'True', 'False'}, 0xF0, 'blob?')

local pf_id = ProtoField.uint8("rtu.SlaveId", "Id", base.DEC_HEX)
local pf_fcode = ProtoField.uint8("rtu.Function", "Function", base.HEX, functions_code)
local pf_size = ProtoField.uint8("rtu.Size", "Size")
local pf_err = ProtoField.uint8("rtu.ErrCode", "ErrCode", base.HEX, exception_code)
local pf_addr = ProtoField.uint16("rtu.RegAddr", "Address")
local pf_cnt = ProtoField.uint16("rtu.RegCount", "Word Count")
local pf_Wcnt = ProtoField.uint16("rtu.WordCount", "Word Count")
local pf_Bcnt = ProtoField.uint8("rtu.ByteCount", "Byte Count")
local pf_bcnt = ProtoField.uint16("rtu.BitCount", "Bit Count")
local pf_reg = ProtoField.uint16("rtu.RegValue", "RegValue", base.HEX)
local pf_crc = ProtoField.uint16("rtu.Crc16", "Crc16", base.HEX)

local pf_bits = ProtoField.uint8("rtu.Bits", "Bits", base.HEX)
local pf_coil = ProtoField.uint8("rtu.Coils", "Coils", base.HEX, coils_code)


g_rtu.fields = {
                pf_id,  pf_fcode, pf_size, pf_err,
                pf_addr,  pf_cnt, pf_Wcnt, pf_Bcnt, pf_bcnt, pf_reg,  pf_crc,
                pf_bits,  pf_coil
                }

-- define dissector function (Tvb, Pinfo, TreeItem)
function g_rtu.dissector(buf, pkt, tree)
    local pktlen = buf:reported_length_remaining()

    if pktlen < 5 then
        return 0
    end

    local id = buf(0, 1):uint()
    local fcode = buf(1, 1):uint()
    local crc = buf(-2, 2)
    local crc16_1 = CRC.crc16plain(buf(0, pktlen-2):tvb())
    local crc16_2 = crc:uint()
    crc16_2 = (crc16_2 >> 8) | (crc16_2 << 8) & 0xFFFF
    
    if crc16_1 ~= crc16_2 then
        return 0
    end

    if not tableContains(functions_code, fcode) then
        return 0
    end

    if (fcode & 0x80) == 0x80 then
        if pktlen == 5 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_err, buf:range(2, 1))
            subtree:add(pf_crc, crc)
        else
            return 0
        end
    elseif fcode == 1 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_bcnt, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        end

        local cnt = buf:range(2, 1):uint()
        if (cnt + 5) == pktlen then
            subtree = tree:add(g_rtu, buf:range(0, pktlen))
            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_Bcnt, buf:range(2, 1))

            local off = 3
            repeat
                subtree:add(pf_bits, buf:range(off, 1))
                off = off + 1
                cnt = cnt - 1
            until cnt == 0
            subtree:add(pf_crc, crc)
        end
    elseif fcode == 2 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_bcnt, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        end

        local cnt = buf:range(2, 1):uint()
        if (cnt + 5) == pktlen then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_Bcnt, buf:range(2, 1))

            local off = 3
            repeat
                subtree:add(pf_bits, buf:range(off, 1))
                off = off + 1
                cnt = cnt - 1
            until cnt == 0

            subtree:add(pf_crc, crc)
        end
    elseif fcode == 3 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_Wcnt, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        else
            local cnt = buf:range(2, 1):uint()
            if (cnt + 5) == pktlen then
                local subtree = tree:add(g_rtu, buf:range(0, pktlen))

                subtree:add(pf_id, buf:range(0, 1))
                subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
                subtree:add(pf_Bcnt, buf:range(2, 1))

                local off = 3
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0

                subtree:add(pf_crc, crc)
            else
                return 0
            end
        end
    elseif fcode == 4 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_Wcnt, buf:range(4, 2))
            subtree:add(pf_crc, buf:range(6, 2))
        else
            local cnt = buf:range(2, 1):uint()
            if (cnt + 5) == pktlen then
                local subtree = tree:add(g_rtu, buf:range(0, pktlen))

                subtree:add(pf_id, buf:range(0, 1))
                subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
                subtree:add(pf_Bcnt, buf:range(2, 1))

                local off = 3
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0

                subtree:add(pf_crc, crc)
            else
                return 0
            end
        end
    elseif fcode == 5 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_coil, buf:range(4, 2))
            subtree:add(pf_crc, crc)

            subtree = tree:add(g_rtu, buf:range(0, pktlen))
            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_coil, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        else
            return 0
        end
    elseif fcode == 6 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_reg, buf:range(4, 2))
            subtree:add(pf_crc, crc)

            subtree = tree:add(g_rtu, buf:range(0, pktlen))
            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(2, 2))
            subtree:add(pf_reg, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        else
            return 0
        end
    elseif fcode == 0x0F then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(2, 2)):append_text(" <Start>")
            subtree:add(pf_bcnt, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        elseif pktlen >= 10 then
            local cnt = buf:range(6, 1):uint()
            if (cnt + 9) == pktlen then
                local subtree = tree:add(g_rtu, buf:range(0, pktlen))

                subtree:add(pf_id, buf:range(0, 1))
                subtree:add(pf_fcode, buf:range(1, 1))
                subtree:add(pf_addr, buf:range(2, 2)):append_text(" <Start>")
                subtree:add(pf_bcnt, buf:range(4, 2))
                subtree:add(pf_Bcnt, buf:range(6, 1))

                local off = 7
                repeat
                    subtree:add(pf_bits, buf:range(off, 1))
                    off = off + 1
                    cnt = cnt - 1
                until cnt == 0

                subtree:add(pf_crc, crc)
            else
                return 0
            end
        else
            return 0
        end
    elseif fcode == 0x10 then
        if pktlen == 8 then
            local subtree = tree:add(g_rtu, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(2, 2)):append_text(" <Start>")
            subtree:add(pf_Wcnt, buf:range(4, 2))
            subtree:add(pf_crc, crc)
        elseif pktlen >= 13 then
            local cnt = buf:range(6, 1):uint()
            if (cnt + 9) == pktlen then
                local subtree = tree:add(g_rtu, buf:range(0, pktlen))

                subtree:add(pf_id, buf:range(0, 1))
                subtree:add(pf_fcode, buf:range(1, 1))
                subtree:add(pf_addr, buf:range(2, 2)):append_text(" <Start>")
                subtree:add(pf_Wcnt, buf:range(4, 2))
                subtree:add(pf_Bcnt, buf:range(6, 1))

                local off = 7
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0

                subtree:add(pf_crc, crc)
            else
                return 0
            end
        else
            return 0
        end
    end


    print("Script ModbusRTU dissector exit")
end

if NotWireshark == nil then
    -- local udp_encap_table = DissectorTable.get("udp.port")

    -- udp_encap_table:add(44444, p_smsProto)
end

print("Script ModbusRTU exit")
